import pandas as pd
import numpy as np
import streamlit as st
from pyecharts import options as opts
from pyecharts.charts import Bar, Line, Pie

# 读取数据集
df = pd.read_csv("dataset.csv")


# 查询指定职位的薪资信息
def query_salary(positionName: str, city: str = None, chart_type: str = None) -> dict:
    print(f"query_salary被调用 {positionName},{city}")
    if not isinstance(positionName, str):
        raise ValueError("positionName 必须是字符串类型")
    if city and not isinstance(city, str):
        raise ValueError("city 必须是字符串类型")

    cond = df['positionName'].str.contains(positionName)
    if city:
        cond &= df['city'] == city
    salaries = df[cond]['mean_salary'].dropna()
    result = {
        "count": len(salaries),
        "average": int(salaries.mean()) if not salaries.empty else 0,
        "positionName": positionName,
        "city": city
    }
    return result


# 不同岗位不同城市学历水平情况
def query_education_level(positionName: str, city: str = None, chart_type: str = None) -> dict:
    print(f"query_education_level被调用 {positionName},{city}")
    if not isinstance(positionName, str):
        raise ValueError("positionName 必须是字符串类型")
    if city and not isinstance(city, str):
        raise ValueError("city 必须是字符串类型")

    cond = df['positionName'].str.contains(positionName)
    if city:
        cond &= df['city'] == city
    education_levels = df[cond]['education'].dropna()
    education_counts = education_levels.value_counts().to_dict()
    result = {
        "education_counts": education_counts,
        "positionName": positionName,
        "city": city
    }
    return result


# 不同岗位不同城市职位数量情况
def query_position_count(positionName: str, city: str = None, chart_type: str = None) -> dict:
    print(f"query_position_count被调用 {positionName},{city}")
    if not isinstance(positionName, str):
        raise ValueError("positionName 必须是字符串类型")
    if city and not isinstance(city, str):
        raise ValueError("city 必须是字符串类型")

    cond = df['positionName'].str.contains(positionName)
    if city:
        cond &= df['city'] == city
    count = len(df[cond])
    result = {
        "count": count,
        "positionName": positionName,
        "city": city
    }
    return result


# 大厂或中小企业不同岗位职位数量情况
def query_position_count_by_company_category(positionName: str, company_category: str, chart_type: str = None) -> dict:
    print(f"query_position_count_by_company_category被调用 {positionName}, {company_category}")
    if not isinstance(positionName, str):
        raise ValueError("positionName 必须是字符串类型")
    if not isinstance(company_category, str):
        raise ValueError("company_category 必须是字符串类型")

    cond = df['positionName'].str.contains(positionName) & df['companySize'].str.contains(company_category)
    count = len(df[cond])
    result = {
        "count": count,
        "positionName": positionName,
        "company_category": company_category
    }
    return result


# 生成Echarts图表数据
def generate_echarts_data(data, chart_type, func_name):
    print(f"generate_echarts_data被调用 {data}, {chart_type}, {func_name}")
    if not data:
        return None, "无数据可绘制图表"

    chart = None
    title = ""

    if chart_type == '柱状图':
        if func_name == 'query_salary':
            chart = (
                Bar()
                .add_xaxis([f"{data['positionName']}({data['city']})" if data['city'] else data['positionName']])
                .add_yaxis("平均薪资", [data["average"]])
                .set_global_opts(title_opts=opts.TitleOpts(title="职位平均薪资柱状图"))
            )
            title = f"{data['positionName']}平均薪资"
        elif func_name == 'query_education_level':
            chart = (
                Bar()
                .add_xaxis(list(data["education_counts"].keys()))
                .add_yaxis("数量", list(data["education_counts"].values()))
                .set_global_opts(title_opts=opts.TitleOpts(title="不同学历水平数量柱状图"))
            )
            title = f"{data['positionName']}学历分布"
        elif func_name == 'query_position_count':
            chart = (
                Bar()
                .add_xaxis([f"{data['positionName']}({data['city']})"])
                .add_yaxis("职位数量", [data["count"]])
                .set_global_opts(title_opts=opts.TitleOpts(title="职位数量柱状图"))
            )
            title = f"{data['positionName']}职位数量"
        elif func_name == 'query_position_count_by_company_category':
            chart = (
                Bar()
                .add_xaxis([f"{data['company_category']} - {data['positionName']}"])
                .add_yaxis("职位数量", [data["count"]])
                .set_global_opts(title_opts=opts.TitleOpts(title="职位数量柱状图"))
            )
            title = f"{data['company_category']}公司{data['positionName']}职位数量"

    elif chart_type == '折线图':
        if func_name == 'query_salary':
            chart = (
                Line()
                .add_xaxis([f"{data['positionName']}({data['city']})" if data['city'] else data['positionName']])
                .add_yaxis("平均薪资", [data["average"]])
                .set_global_opts(title_opts=opts.TitleOpts(title="职位平均薪资折线图"))
            )
            title = f"{data['positionName']}平均薪资"
        elif func_name == 'query_education_level':
            chart = (
                Line()
                .add_xaxis(list(data["education_counts"].keys()))
                .add_yaxis("数量", list(data["education_counts"].values()))
                .set_global_opts(title_opts=opts.TitleOpts(title="不同学历水平数量折线图"))
            )
            title = f"{data['positionName']}学历分布"
        elif func_name == 'query_position_count':
            chart = (
                Line()
                .add_xaxis([f"{data['positionName']}({data['city']})"])
                .add_yaxis("职位数量", [data["count"]])
                .set_global_opts(title_opts=opts.TitleOpts(title="职位数量折线图"))
            )
            title = f"{data['positionName']}职位数量"
        elif func_name == 'query_position_count_by_company_category':
            chart = (
                Line()
                .add_xaxis([f"{data['company_category']} - {data['positionName']}"])
                .add_yaxis("职位数量", [data["count"]])
                .set_global_opts(title_opts=opts.TitleOpts(title="职位数量折线图"))
            )
            title = f"{data['company_category']}公司{data['positionName']}职位数量"

    elif chart_type == '饼图':
        if func_name == 'query_salary':
            chart = (
                Pie()
                .add(
                    series_name="平均薪资",
                    data_pair=[(f"{data['positionName']}({data['city']})" if data['city'] else data['positionName'], data["average"])],
                    radius=["40%", "75%"],
                )
                .set_global_opts(title_opts=opts.TitleOpts(title="职位平均薪资饼图"))
            )
            title = f"{data['positionName']}平均薪资"
        elif func_name == 'query_education_level':
            chart = (
                Pie()
                .add(
                    series_name="学历水平",
                    data_pair=[(k, v) for k, v in data["education_counts"].items()],
                    radius=["40%", "75%"],
                )
                .set_global_opts(title_opts=opts.TitleOpts(title="不同学历水平数量饼图"))
            )
            title = f"{data['positionName']}学历分布"
        elif func_name == 'query_position_count':
            chart = (
                Pie()
                .add(
                    series_name="职位数量",
                    data_pair=[(f"{data['positionName']}({data['city']})", data["count"])],
                    radius=["40%", "75%"],
                )
                .set_global_opts(title_opts=opts.TitleOpts(title="职位数量饼图"))
            )
            title = f"{data['positionName']}职位数量"
        else:
            return None, "不支持生成该类型Echarts图像"
    else:
        return None, "不支持生成该类型Echarts图像"

    return chart, title


# 定义工具列表
tools = [
    {
        "type": "function",
        "function": {
            "name": "query_salary",
            "description": "查询指定职位的薪资信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "positionName": {"type": "string", "description": "职位名称，如：算法工程师"},
                    "city": {"type": "string", "description": "城市名称，如：北京"}
                },
                "required": ["positionName"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_education_level",
            "description": "查询不同岗位不同城市学历水平情况",
            "parameters": {
                "type": "object",
                "properties": {
                    "positionName": {"type": "string", "description": "职位名称，如：算法工程师"},
                    "city": {"type": "string", "description": "城市名称，如：北京"}
                },
                "required": ["positionName"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_position_count",
            "description": "查询不同岗位不同城市职位数量情况",
            "parameters": {
                "type": "object",
                "properties": {
                    "positionName": {"type": "string", "description": "职位名称，如：算法工程师"},
                    "city": {"type": "string", "description": "城市名称，如：北京"}
                },
                "required": ["positionName"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_position_count_by_company_category",
            "description": "查询大厂或中小企业不同岗位职位数量情况",
            "parameters": {
                "type": "object",
                "properties": {
                    "positionName": {"type": "string", "description": "职位名称，如：算法工程师"},
                    "company_category": {"type": "string", "description": "公司规模类别，如：150-500人"}
                },
                "required": ["positionName", "company_category"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "generate_echarts_data",
            "description": "基于查询结果生成 ECharts 图表数据",
            "parameters": {
                "type": "object",
                "properties": {
                    "data": {"type": "object", "description": "前一步骤返回的结果数据"},
                    "chart_type": {"type": "string", "description": "图表类型，如：柱状图、折线图、饼图"},
                    "func_name": {"type": "string", "description": "前一步骤调用的函数名称"}
                },
                "required": ["data", "chart_type", "func_name"]
            }
        }
    }
]